import pod5
import argparse
import os
import pysam

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--pod5_dir', type=str, required=False,
                        help="id")
    parser.add_argument('--input', type=str, required=False,
                        help="id")
    parser.add_argument('--extra', type=str, required=False,
                        help="id")
    parser.add_argument('--bam', type=str, required=False,
                        help="id")
    parser.add_argument('--output', type=str, required=False,
                        help="id")
    return parser.parse_args()

def read_pod5_address(input,extra):
    inputs = []
    with open(input,'r') as rb:
        for line in rb:
            last_slash_index = line.rfind('/')
            pod5_name=line[last_slash_index + 1:]
            if pod5_name not in extra:
                inputs.append(line)#os.path.abspath(line)
    return inputs

def read_extra(input):
    inputs = []
    with open(input,'r') as rb:
        for line in rb:
            last_slash_index = line.rfind('/')
            extracted_string = line[last_slash_index + 1:]
            inputs.append(extracted_string)
    return inputs

def read_bam(inbam):
    inputs = []
    input_bam = pysam.AlignmentFile(inbam, 'rb',check_sq=False)
    for read in input_bam:
        inputs.append(read.query_name)
    return inputs

if __name__ == '__main__':
    args=parse_args()
    writer = open(args.output, 'w')
    if args.bam is None:
        extra=read_extra(args.extra)
        inputs=read_pod5_address(args.input,extra)
        for pod5_name in inputs:
            with pod5.Reader(str(pod5_name)) as reader:
                for read_record in reader.reads():
                    writer.write(str(read_record.read_id)+'\n')
    else:
        inputs=read_bam(args.bam)
        for read_id in inputs:
            writer.write(str(read_id)+'\n')
    
    
    